from iotdb.Session import Session
from iotdb.utils.IoTDBConstants import TSDataType, TSEncoding, Compressor
from iotdb.utils.Tablet import Tablet
import time
from datetime import datetime
import pandas as pd
import numpy as np
from tqdm import tqdm

np.object = object
"""
    IotDB_utils.py

    @description:
        IotDB工具类
    @author:
        zhenyu
    @date:
        2024/12/31
"""


class IotDB_utils:
    tags_lst = []
    ts_path_lst = []
    measurements = []
    data_type_lst = []

    def __init__(self, host, port, user, pwd):
        self.host = str(host)
        self.port = int(port)
        self.user = str(user)
        self.password = str(pwd)
        self.maxSize = 100

        self.device = ""
        self.device_ids = device + ""
        self.storage_has_create = False

        self.session = Session(self.host, self.port, self.user, self.password, self.maxSize)
        self.session.open(enable_rpc_compression=False)
        print("登录成功")

    def create_ts(self, ts_path_lst, data_type_lst, tags_lst):
        encoding_lst = [TSEncoding.GORILLA for _ in range(len(data_type_lst))]
        compressor_lst = [Compressor.SNAPPY for _ in range(len(data_type_lst))]
        self.session.create_multi_time_series(
            ts_path_lst=ts_path_lst,
            data_type_lst=data_type_lst,
            encoding_lst=encoding_lst,
            compressor_lst=compressor_lst,
            tags_lst=tags_lst
        )

    def __set_ts_path_and_measurement_lst(self, col, device_id):
        if col in ["wind_speed", "wind_angle", "air_pressure"]:
            self.ts_path_lst.append(device_id + "." + "environment" + "." + col)
            self.measurements.append("environment" + "." + col)
        elif col in ["battery_voltage", "battery_current", "battery_remind"]:
            self.ts_path_lst.append(device_id + "." + "_power" + "." + col)
            self.measurements.append("_power" + "." + col)
        elif col in ["gps_x", "gps_y", "gps_z", "real_lat", "real_long", "aim_lat", "aim_long", "aim_z"]:
            self.ts_path_lst.append(device_id + "." + "position" + "." + col)
            self.measurements.append("position" + "." + col)
        elif col in ["o_x", "o_y", "o_z", "o_w", "v_x", "v_y", "v_z", "la_x", "la_y", "la_z"]:
            self.ts_path_lst.append(device_id + "." + "attitude" + "." + col)
            self.measurements.append("attitude" + "." + col)
        else:
            self.ts_path_lst.append(device_id + "." + "other." + col)
            self.measurements.append("other" + "." + col)
        # self.measurements.append(col)

    def __set_tags_lst(self, col):
        if col == 'wind_speed' or col == 'wind_angle' or col == 'air_pressure':
            self.tags_lst.append({"dim": "stability+robustness+perception"})
        elif col == 'battery_voltage' or col == 'battery_current' or col == 'battery_remain':
            self.tags_lst.append({"dim": "robustness"})
        elif col == 'gps_x' or col == 'gps_y' or col == 'gps_z':
            self.tags_lst.append({"dim": "stability+perception"})
        elif col == 'real_lat' or col == 'real_long' or col == 'aim_lat' or col == 'aim_long' or col == 'aim_z':
            self.tags_lst.append({"dim": "perception"})
        elif col == 'o_x' or col == 'o_y' or col == 'o_z' or col == 'o_w':
            self.tags_lst.append({"dim": "flexibility+perception"})
        elif col == 'v_x' or col == 'v_y' or col == 'v_z' or col == 'la_x' or col == 'la_y' or col == 'la_z':
            self.tags_lst.append({"dim": "flexibility"})
        else:
            self.tags_lst.append({"dim": "other"})

    def fill_dataframe_nan(self, df, fill_method='mean', fill_value=0):
        """
        处理 DataFrame中的空值。
        """

        # 处理 df 中的空值
        if fill_method == 'mean':
            df = df.fillna(df.mean())
        elif fill_method == 'median':
            df = df.fillna(df.median())
        elif fill_method == 'mode':
            df = df.fillna(df.mode().iloc[0])
        elif fill_method == 'constant':
            df = df.fillna(fill_value)
        else:
            raise ValueError(f"Unsupported fill method: {fill_method}")
        return df

    def IotDB_insert(self, excel_path: str, device: str, device_id: str):
        if not self.storage_has_create:
            self.session.set_storage_group(device)
            self.storage_has_create = True

        # 使用pandas的read_excel函数读取Excel文件
        df = pd.read_csv(excel_path)
        df = self.fill_dataframe_nan(df, fill_method='constant')

        for col in df.columns:
            # 设置
            self.__set_ts_path_and_measurement_lst(col=col, device_id=device_id)
            self.data_type_lst.append(TSDataType.DOUBLE)
            self.__set_tags_lst(col)

        print(self.ts_path_lst)
        print(self.measurements)
        print(self.tags_lst)

        # 设置时间序列
        if not self.session.check_time_series_exists(device_id):
            encoding_lst = [TSEncoding.GORILLA for _ in range(len(self.data_type_lst))]
            compressor_lst = [Compressor.SNAPPY for _ in range(len(self.data_type_lst))]
            self.session.create_multi_time_series(
                ts_path_lst=self.ts_path_lst,
                data_type_lst=self.data_type_lst,
                encoding_lst=encoding_lst,
                compressor_lst=compressor_lst,
                tags_lst=self.tags_lst
            )
            print("时间序列写入成功")

        data = {}
        data_nums = 0
        for col in df.columns:
            data[self.ts_path_lst[data_nums]] = np.array(df[col].tolist())
            data_nums += 1

        df_input = pd.DataFrame(data)

        # 设置起始时间
        new_time = datetime.now().replace(month=10, day=8, hour=14)
        start_time = int(datetime.timestamp(new_time) * 1000)

        # 按20ms插入数据
        for index, row in tqdm(df_input.iterrows(), total=len(df_input), desc="insert"):
            for col_index, column in enumerate(df_input.columns):
                parts = column.rsplit('.', 1)
                device_id = parts[0]
                measurements = parts[1]
                self.session.insert_record(device_id, start_time, [measurements], [self.data_type_lst[col_index]],
                                           [row[column]])

            # 20ms
            start_time += 20
            time.sleep(0.02)


if __name__ == '__main__':
    device = "root.fafs"
    device_id = device + ".uavy"
    excel_file = "./data/UavY_P0A10S2_1.csv"

    username_ = "root"
    password_ = "root"
    ip = "127.0.0.1"
    port_ = 6667
    maxSize = 100

    iotdb_test = IotDB_utils(ip, port_, username_, password_)
    iotdb_test.IotDB_insert(excel_file, device, device_id)
